package com.demo.LSTM;

import org.deeplearning4j.datasets.iterator.impl.ListDataSetIterator;
import org.deeplearning4j.nn.api.Layer;
import org.deeplearning4j.nn.conf.*;
import org.deeplearning4j.nn.conf.layers.DenseLayer;
import org.deeplearning4j.nn.conf.layers.LSTM;
import org.deeplearning4j.nn.conf.layers.RnnOutputLayer;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.deeplearning4j.nn.weights.WeightInit;
import org.nd4j.evaluation.regression.RegressionEvaluation;
import org.nd4j.linalg.activations.Activation;
import org.nd4j.linalg.api.buffer.DataType;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.dataset.DataSet;
import org.nd4j.linalg.factory.Nd4j;
import org.nd4j.linalg.learning.config.Adam;
import org.nd4j.linalg.lossfunctions.LossFunctions;

import java.util.ArrayList;
import java.util.List;

public class LSTMLoadPredictionDemoWithFactors {
    public static void main(String[] args) {
        int numInputs = 8; // 负荷和各种影响因素数据
        int numOutputs = 1;
        int numHiddenUnits = 10;
        int numEpochs = 100;

        // 创建数据集
        double[] loadArray = {0.1, 0.2, 0.3, 0.4};
        double[] seasonArray = {0.5, 0.6, 0.7, 0.8};
        double[] holidayArray = {0.0, 1.0, 0.0, 1.0};
        double[] loadCharacteristicsArray = {0.9, 1.0, 0.8, 0.6};
        double[] weatherArray = {25.0, 26.0, 27.0, 28.0};
        double[] policyArray = {1.0, 1.0, 0.0, 0.0};
        double[] economicArray = {0.8, 0.9, 1.1, 1.2};
        double[] loadManagementArray = {1.2, 1.1, 1.0, 0.9};

        List<DataSet> dataList = new ArrayList<>();
        for (int i = 0; i < loadArray.length; i++) {
            double[] input = {loadArray[i], seasonArray[i], holidayArray[i], loadCharacteristicsArray[i], weatherArray[i], policyArray[i], economicArray[i], loadManagementArray[i]};
            double[] output = {loadArray[i]};
            dataList.add(new DataSet(Nd4j.create(input, new int[]{1, numInputs, 1}), Nd4j.create(output, new int[]{1, numOutputs, 1})));
        }

        ListDataSetIterator<DataSet> iterator = new ListDataSetIterator<>(dataList, 1);

        // 构建模型
        MultiLayerConfiguration conf = new NeuralNetConfiguration.Builder()
                .seed(123)
                .weightInit(WeightInit.XAVIER)
                .updater(new Adam(0.01))
                .list()
                .layer(new LSTM.Builder().nIn(numInputs).nOut(numHiddenUnits).activation(Activation.TANH).build())
                .layer(new RnnOutputLayer.Builder().nIn(numHiddenUnits).nOut(numOutputs).lossFunction(LossFunctions.LossFunction.MSE).activation(Activation.IDENTITY).build())
                .build();

        MultiLayerNetwork model = new MultiLayerNetwork(conf);
        model.init();

        // 训练模型
        for (int i = 0; i < numEpochs; i++) {
            model.fit(iterator);
        }

        // 准备未来4小时内的输入数据
        double[] futureInput = {0.5, 0.9, 1.0, 0.7, 30.0, 1.0, 1.2, 1.1};

        // 进行预测
        INDArray input = Nd4j.create(futureInput, new int[]{1, numInputs, 1});
        INDArray predictedOutput = model.rnnTimeStep(input);

        System.out.println(predictedOutput.shapeInfoToString());

        // 解析预测结果
        float[] doubles = predictedOutput.toFloatVector();

        System.out.println("未来4小时内的负荷预测结果：");
        for (float prediction : doubles) {
            System.out.println(prediction);
        }
    }
}
